package org.shiqianghh.tanslate;

import java.util.Map;

import com.google.common.collect.Maps;

import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;

public class TacotronTranslator implements Translator<NDList, Map> {

    public TacotronTranslator() {

    }
    /** {@inheritDoc} */
    @Override
    public NDList processInput(TranslatorContext ctx, NDList input) {  
        return input;
    }

    /** {@inheritDoc} */
    @Override
    public Map<String,NDArray> processOutput(TranslatorContext ctx, NDList list) { 
    	Map<String,NDArray> re = Maps.newHashMap();
    	NDArray mels = list.get(0);
    	NDArray mels_postnet = list.get(1);
    	NDArray gates = list.get(2);
    	NDArray alignments = list.get(3);
    	mels.detach();
    	mels_postnet.detach();
    	gates.detach();
    	alignments.detach(); 
    	re.put("mels", mels);
    	re.put("mels_postnet", mels_postnet);
    	re.put("gates", gates);
    	re.put("alignments", alignments); 
        return re;
    }

    /** {@inheritDoc} */
    @Override
    public Batchifier getBatchifier() {
        return Batchifier.STACK;
    }

}
